from http import HTTPStatus
import dashscope
import os
import requests
from requests_toolbelt.multipart.encoder import MultipartEncoder
from dataclasses import dataclass

CWD = os.path.dirname(os.path.abspath(__file__))
API_KEY = os.environ['DASHSCOPE_API_KEY']
ALIYUN_DOMAIN = "https://dashscope.aliyuncs.com/api/v1"
FILE_URL = f"{ALIYUN_DOMAIN}/files"
TRAIN_FACECHAIN_URL = f"{ALIYUN_DOMAIN}/fine-tunes"

def upload_file(file_path):
    print(f"pic: {file_path}")
    data = {
        "descriptions": "test2",
    }
    headers = {
    #    "content-type": "multipart/form-data",
       "Authorization": f"Bearer {API_KEY}"
    }
    files = {
        # "descriptions": "test2",
        ("files", open(file_path, 'rb')),
    }
    result = requests.post(url=FILE_URL, params=data, headers=headers, files=files)
    print(result.json())

    file_ids = []
    if result.status_code == 200:
        files_body = result.json()['data']['uploaded_files']
        for file in files_body:
            file_ids.append(file['file_id'])
    print(file_ids)
    return file_ids

def check_uploaded_files():
    print("Check uploaded files ...")
    headers = {
        "content-type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    result = requests.get(url=FILE_URL, headers=headers)
    print(result.json())
    file_ids = []
    if result.status_code == 200:
        data = result.json()['data']
        for file in data['files']:
            file_ids.append(file['file_id'])
    return file_ids


def delete_uploaded_files(file_ids):
    print("Delete uploaded files ...")
    for file_id in file_ids:
        url = f"{FILE_URL}/{file_id}"
        headers = {
            "content-type": "application/json",
            "Authorization": f"Bearer {API_KEY}"
        }
        response = requests.delete(url=url, headers=headers)
        print(response.json())


def train_facechain_model(train_file_id):
    print(f'Starting to finetune with file: {train_file_id}')
    headers = {
        "content-type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }

    data = {
        "model": "facechain-finetune",
        "training_file_ids": [train_file_id],
    }

    result = requests.post(url=TRAIN_FACECHAIN_URL, json=data, headers=headers)
    print(result.json())
    if result.status_code == 200:
        output = result.json()['output']
        job_id = output['job_id']
        status = output['status']
        return job_id, status

def finetune_image_generator(resource_id, picture_style):
    url = f"{ALIYUN_DOMAIN}/services/aigc/album/gen_potrait"
    print(f'Starting to generate image ...')
    headers = {
        "content-type": "application/json",
        "Authorization": f"Bearer {API_KEY}",
        "X-DashScope-Async": "enable",
    }
    data = {
        "model": "facechain-generation",
        "resources": [
            {
                "resource_id": resource_id,
                "resource_type": "facelora",
            },
        ],
        "parameters": {
            "style": picture_style,
            "size": "768*1024",
            "n": 4,
        }
    }
    result = requests.post(url=url, json=data, headers=headers)
    print(result.json())
    if result.status_code == 200:
        return result.json()['output']['task_id']

def check_image_ready(task_id):
    print(f'Check if the image is generated ...')
    if task_id is None:
        print(f'model training is not ready ...')
        return

    url = f"{ALIYUN_DOMAIN}/tasks/{task_id}"
    headers = {
        "content-type": "application/json",
        "Authorization": f"Bearer {API_KEY}",
    }
    result = requests.get(url=url, headers=headers)
    print(result.json())
    picture_urls = []
    if result.status_code == 200:
        output = result.json()['output'] 
        task_status = output['task_status']
        if task_status == "SUCCEEDED":
            for url in output['results']:
                picture_urls.append(url['url'])
    return picture_urls 

def check_finetune_image(*image_urls):
    url = f"{ALIYUN_DOMAIN}/services/vision/facedetection/detect"
    headers = {
        "content-type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    data = {
        "model": "facechain-facedetect",
        "input": {
            "images": image_urls,
            # "images": ["http://xxx.jpg",]
        }
    }
    status = requests.post(url=url, json=data, headers=headers)
    print(status.json())


def check_train_status(job_id):
    url = f"{TRAIN_FACECHAIN_URL}/{job_id}"
    headers = {
        "content-type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    result = requests.get(url=url, headers=headers)
    print(result.json())
    if result.status_code == 200:
        output = result.json()['output'] 
        if output['status'] == "SUCCEEDED":
            return output['finetuned_output']

def download_images(image_urls):
    print("Start to download pictures ...")
    for url in image_urls:
        result = requests.get(url=url, stream=True)
        pic_name = url.split("Signature=")[-1]
        pic_path = os.path.join(CWD, f"{pic_name}.png")
        print(pic_path)
        with open(pic_path, 'wb') as f:
            for chunk in result.iter_content(chunk_size=10*1024):
                f.write(chunk)


@dataclass
class PictureStyle:
    business_male = "f_business_male"
    idcard_male = "f_idcard_male"
    idcard_female = "f_idcard_female"
    business_female = "f_business_female"
    springflower_female = "m_springflower_female"
    summersport_female = "f_summersport_female"
    autumnleaf_female = "f_autumnleaf_female"
    winterchinese_female = "m_winterchinese_female"
    hongkongvintage_female = "f_hongkongvintage_female"
    lightportray_female = "f_lightportray_female"


if __name__ == '__main__':
    # Comment all lines and uncomment them one by one to check the result
    pic_1 = os.path.join(CWD, 'p731277.zip')
    file_ids = upload_file(pic_1)
    file_ids = check_uploaded_files()
    print(f"Uploaded files: {file_ids}")
    # file_ids = ['46728968-1376-4fc5-b7ec-ea17ac664b63']
    delete_uploaded_files(file_ids)
    job_id = train_facechain_model(file_ids[0])
    # job_id = 'ft-202312062031-26c3'
    trained_model = check_train_status(job_id)
    print(f"trained model: {trained_model}")
    task_id = finetune_image_generator(resource_id=trained_model, picture_style=PictureStyle.autumnleaf_female)
    # task_id = "ce7b0d60-4fca-4509-b000-b600254b2fbb"
    picture_urls = check_image_ready(task_id)
    print(f"picture urls: {picture_urls}")
    download_images(picture_urls)
